package com.alison.datastream.chapter11_connector;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.common.typeutils.base.StringSerializer;
import org.apache.flink.api.common.typeutils.base.VoidSerializer;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.RestOptions;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStream;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.TwoPhaseCommitSinkFunction;
import org.apache.flink.streaming.api.functions.source.RichSourceFunction;

import java.io.BufferedWriter;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.LocalDateTime;
import java.time.format.DateTimeFormatter;

public class E3_TransactionWriteSinkExample {

    public static void main(String[] args) throws Exception {
        Configuration conf = new Configuration();
        // 访问 http://localhost:8082 可以看到Flink Web UI
        conf.setInteger(RestOptions.PORT, 8082);
        // 创建本地执行环境，并行度为1
        StreamExecutionEnvironment env = StreamExecutionEnvironment.createLocalEnvironment(1, conf);
        // 每隔5秒进行一次Checkpoint
        env.getCheckpointConfig().setCheckpointInterval(5 * 1000);

        DataStream<Tuple2<String, Integer>> countStream = env.addSource(new CheckpointedSource());
        // 每隔一定时间模拟一次失败
        DataStream<Tuple2<String, Integer>> result = countStream.map(new FailingMapper(20));

        // 类Unix系统的临时文件夹在/tmp下
        // Windows用户需要修改这个目录
        String preCommitPath = "/tmp/flink-sink-precommit";
        String commitedPath = "/tmp/flink-sink-commited";

        if (!Files.exists(Paths.get(preCommitPath))) {
            Files.createDirectory(Paths.get(preCommitPath));
        }
        if (!Files.exists(Paths.get(commitedPath))) {
            Files.createDirectory(Paths.get(commitedPath));
        }
        // 使用Exactly-Once语义的Sink，执行本程序时可以查看相应的输出目录，查看数据
        result.addSink(new TwoPhaseFileSink(preCommitPath, commitedPath));
        // 数据打印到屏幕上，无Exactly-Once保障，有数据重发现象
        result.print();
        env.execute("two file sink");
    }
    public static class FailingMapper
            implements MapFunction<Tuple2<String, Integer>, Tuple2<String, Integer>> {

        private int count = 0;
        private int failInterval;

        public FailingMapper(int failInterval) {
            this.failInterval = failInterval;
        }

        @Override
        public Tuple2<String, Integer> map(Tuple2<String, Integer> in) {
            count += 1;
            if (count > failInterval) {
//                throw new RuntimeException("job fail! show how flink checkpoint and recovery");
            }
            return in;
        }
    }
    public static class CheckpointedSource
            extends RichSourceFunction<Tuple2<String, Integer>>
            implements CheckpointedFunction {

        private int offset;
        private boolean isRunning = true;
        private ListState<Integer> offsetState;

        @Override
        public void run(SourceContext<Tuple2<String, Integer>> ctx) throws Exception {
            while (isRunning) {
                Thread.sleep(100);
                synchronized (ctx.getCheckpointLock()) {
                    ctx.collect(new Tuple2<>("" + offset, 1));
                    offset++;
                }
                if (offset == 1000) {
                    isRunning = false;
                }
            }
        }

        @Override
        public void cancel() {
            isRunning = false;
        }

        @Override
        public void snapshotState(FunctionSnapshotContext snapshotContext) throws Exception {
            // 清除上次状态
            offsetState.clear();
            // 将最新的offset添加到状态中
            offsetState.add(offset);
        }

        @Override
        public void initializeState(FunctionInitializationContext initializationContext) throws Exception {
            // 初始化offsetState
            ListStateDescriptor<Integer> desc = new ListStateDescriptor<Integer>("offset", Types.INT);
            offsetState = initializationContext.getOperatorStateStore().getListState(desc);

            Iterable<Integer> iter = offsetState.get();
            if (iter == null || !iter.iterator().hasNext()) {
                // 第一次初始化，从0开始计数
                offset = 0;
            } else {
                // 从状态中恢复offset
                offset = iter.iterator().next();
            }
        }
    }

    public static class TwoPhaseFileSink extends TwoPhaseCommitSinkFunction<Tuple2<String, Integer>, String, Void> {
        // 缓存
        private BufferedWriter transactionWriter;
        private String preCommitPath;
        private String commitedPath;

        public TwoPhaseFileSink(String preCommitPath, String commitedPath) {
            super(StringSerializer.INSTANCE, VoidSerializer.INSTANCE);
            this.preCommitPath = preCommitPath;
            this.commitedPath = commitedPath;
        }

        @Override
        public void invoke(String transaction, Tuple2<String, Integer> in, Context context) throws Exception {
            transactionWriter.write(in.f0 + " " + in.f1 + "\n");
        }

        @Override
        public String beginTransaction() throws Exception {
            String time = LocalDateTime.now().format(DateTimeFormatter.ISO_LOCAL_DATE_TIME);
            int subTaskIdx = getRuntimeContext().getIndexOfThisSubtask();
            String fileName = time + "-" + subTaskIdx;
            Path preCommitFilePath = Paths.get(preCommitPath + "/" + fileName.replace(":", "-"));
            // 创建一个存储本次事务的文件
            Files.createFile(preCommitFilePath);
            transactionWriter = Files.newBufferedWriter(preCommitFilePath);
            System.out.println("Transaction File: " + preCommitFilePath);

            return fileName;
        }

        @Override
        public void preCommit(String transaction) throws Exception {
            // 将当前数据由内存写入磁盘
            transactionWriter.flush();
            transactionWriter.close();
        }

        @Override
        public void commit(String transaction) {
            Path preCommitFilePath = Paths.get(preCommitPath + "/" + transaction.replace(":", "-"));
            if (Files.exists(preCommitFilePath)) {
                Path commitedFilePath = Paths.get(commitedPath + "/" + transaction.replace(":", "-"));
                try {
                    Files.move(preCommitFilePath, commitedFilePath);
                } catch (Exception e) {
                    System.out.println(e.getMessage());
                }
            }
        }

        @Override
        public void abort(String transaction) {
            Path preCommitFilePath = Paths.get(preCommitPath + "/" + transaction);

            // 如果中途遇到中断，将文件删除
            if (Files.exists(preCommitFilePath)) {
                try {
                    Files.delete(preCommitFilePath);
                } catch (Exception e) {
                    System.out.println(e.getMessage());
                }
            }
        }
    }
}
